package midend.IR_Optimizer;

import Utils.BackEdge;
import Utils.Loop;
import midend.IR_Value.*;
import midend.IR_Instruction.*;
import midend.IR_Type.*;
import java.util.*;

public class LoopInvariantCodeMotion {

    private Map<IR_BaseBlock, Set<IR_BaseBlock>> dominators;
    private Map<IR_BaseBlock, IR_BaseBlock> immediateDominators;
    private List<BackEdge> backEdges;
    private List<Loop> loops;
    private IR_FuncBlock curfuncBlock;

    public LoopInvariantCodeMotion() {
        this.dominators = new HashMap<>();
        this.immediateDominators = new HashMap<>();
        this.backEdges = new ArrayList<>();
        this.loops = new ArrayList<>();
    }

    /**
     * 主要的循环不变式外提优化入口
     */
    public void optimizeFunction(IR_FuncBlock function) {
        curfuncBlock = function;
        // 1. 构建CFG
        //function.get_CFG2();

        // 2. 计算支配关系
        computeDominators(function);

        // 3. 识别回边
        identifyBackEdges(function);

        // 4. 识别自然循环
        identifyNaturalLoops(function);

        // 5. 对每个循环进行不变式外提
        for (Loop loop : loops) {
            performLoopInvariantCodeMotion(loop);
        }
    }

    /**
     * 计算支配关系
     */
    private void computeDominators(IR_FuncBlock function) {
        List<IR_BaseBlock> blocks = new ArrayList<>(function.get_bb_List());
        if (blocks.isEmpty()) return;

        IR_BaseBlock entry = blocks.get(0);
        //System.out.println(entry);
        //System.out.println(blocks);
        // 初始化：入口块只支配自己，其他块被所有块支配
        for (IR_BaseBlock block : blocks) {
            dominators.put(block, new HashSet<>());
            if (block == entry) {
                dominators.get(block).add(block);
            } else {
                dominators.get(block).addAll(blocks);
            }
        }

        // 迭代计算支配关系
        boolean changed = true;
        while (changed) {
            changed = false;
            for (IR_BaseBlock block : blocks) {
                if (block == entry) continue;

                Set<IR_BaseBlock> newDom = new HashSet<>(blocks);

                // 计算所有前驱的支配集合的交集
                for (IR_BaseBlock pred : block.getPreBlock()) {
                    newDom.retainAll(dominators.get(pred));
                }

                // 添加自己
                newDom.add(block);

                if (!newDom.equals(dominators.get(block))) {
                    dominators.put(block, newDom);
                    changed = true;
                }
            }
        }

        // 计算直接支配关系
        computeImmediateDominators(blocks);
    }

    /**
     * 计算直接支配关系
     */
    private void computeImmediateDominators(List<IR_BaseBlock> blocks) {
        for (IR_BaseBlock block : blocks) {
            Set<IR_BaseBlock> doms = new HashSet<>(dominators.get(block));
            doms.remove(block); // 移除自己

            IR_BaseBlock idom = null;
            for (IR_BaseBlock d : doms) {
                boolean isImmediate = true;
                for (IR_BaseBlock other : doms) {
                    if (other != d && dominators.get(other).contains(d)) {
                        isImmediate = false;
                        break;
                    }
                }
                if (isImmediate) {
                    idom = d;
                    break;
                }
            }

            if (idom != null) {
                immediateDominators.put(block, idom);
            }
        }
    }

    /**
     * 识别回边
     */
    private void identifyBackEdges(IR_FuncBlock function) {
        backEdges.clear();

        for (IR_BaseBlock block : function.get_bb_List()) {
            for (IR_BaseBlock succ : block.getNextBlock()) {
                // 如果后继支配当前块，则是回边
                if (dominators.get(block).contains(succ)) {
                    backEdges.add(new BackEdge(block, succ));
                }
            }
        }
    }

    /**
     * 识别自然循环
     */
    private void identifyNaturalLoops(IR_FuncBlock function) {
        loops.clear();

        for (BackEdge backEdge : backEdges) {
            Loop loop = new Loop(backEdge.to);

            // 使用深度优先搜索找到循环体
            Set<IR_BaseBlock> visited = new HashSet<>();
            Stack<IR_BaseBlock> stack = new Stack<>();

            stack.push(backEdge.from);
            loop.body.add(backEdge.to);

            while (!stack.isEmpty()) {
                IR_BaseBlock current = stack.pop();

                if (visited.contains(current) || current == backEdge.to) {
                    continue;
                }

                visited.add(current);
                loop.body.add(current);

                // 添加所有前驱
                for (IR_BaseBlock pred : current.getPreBlock()) {
                    if (!visited.contains(pred)) {
                        stack.push(pred);
                    }
                }
            }

            // 找到循环出口
            for (IR_BaseBlock bodyBlock : loop.body) {
                for (IR_BaseBlock succ : bodyBlock.getNextBlock()) {
                    if (!loop.body.contains(succ)) {
                        loop.exits.add(succ);
                    }
                }
            }

            loops.add(loop);
        }
    }

    /**
     * 执行循环不变式外提
     */
    private void performLoopInvariantCodeMotion(Loop loop) {
        // 1. 创建预头部
        createPreHeader(loop);

        // 2. 识别循环不变式
        identifyLoopInvariants(loop);

        // 3. 外提循环不变式
        hoistInvariants(loop);
    }

    /**
     * 创建预头部基本块
     */
    private void createPreHeader(Loop loop) {
        IR_BaseBlock header = loop.header;
        IR_FuncBlock function = curfuncBlock;

        // 创建预头部
        IR_BaseBlock preHeader = new IR_BaseBlock("preheader_" + header.getName(), function);
        loop.preHeader = preHeader;

        // 添加到函数的基本块列表
        function.get_bb_List().add(preHeader);

        // 重新连接CFG
        List<IR_BaseBlock> outsidePredsOfHeader = new ArrayList<>();
        outsidePredsOfHeader.addAll(header.getPreBlock());

//      //针对存在continue的情况下，此时也就是说一个循环头有两个loop
        List<IR_BaseBlock> deletelist = new ArrayList<>();
        for(Loop tmploop : loops) {
            if (!tmploop.header.equals(loop.header)) {
                continue;
            }
            //排除从循环内过来的
            for (IR_BaseBlock block : outsidePredsOfHeader) {
                if (tmploop.body.contains(block) && !deletelist.contains(block)) {
                    deletelist.add(block);
                }
            }
        }

        for(IR_BaseBlock bb:deletelist){
            outsidePredsOfHeader.remove(bb);
        }

        // 将外部前驱连接到预头部
        for (IR_BaseBlock pred : outsidePredsOfHeader) {
            pred.getNextBlock().remove(header);
            pred.getNextBlock().add(preHeader);
            preHeader.getPreBlock().add(pred);

            // 更新分支指令
            updateBranchTarget(pred, header, preHeader);

            //修改可能由于前驱后继导致的 phi指令错误
            updatePhiTarget(pred, header, preHeader);
        }

        //updatePhiTarget(pred, header, preHeader);

        // 预头部连接到循环头
        preHeader.getNextBlock().add(header);
        header.getPreBlock().removeAll(outsidePredsOfHeader);
        header.getPreBlock().add(preHeader);

        // 添加无条件分支到循环头
        Branch_Instruction branch = new Branch_Instruction(Arrays.asList(header), preHeader);
        preHeader.add_InstructionList(branch);
    }

    /**
     * 识别循环不变式
     */
    private void identifyLoopInvariants(Loop loop) {
        Set<Instruction> invariants = new HashSet<>();
        boolean changed = true;

        while (changed) {
            changed = false;

            for (IR_BaseBlock block : loop.body) {
                for (Instruction inst : block.get_Instruction_List()) {
                    if (invariants.contains(inst)) continue;

                    // 检查是否是循环不变式
                    if (isLoopInvariant(inst, loop, invariants)) {
                        invariants.add(inst);
                        changed = true;
                    }
                }
            }
        }

        loop.invariants = invariants;
    }

    /**
     * 判断指令是否是循环不变式
     */
    private boolean isLoopInvariant(Instruction inst, Loop loop, Set<Instruction> knownInvariants) {
        // 跳过某些不能移动的指令
        if (inst instanceof Store_Instruction ||
                inst instanceof Call_Instruction ||
                inst instanceof Branch_Instruction ||
                inst instanceof Ret_Instruction ||
                inst instanceof Phi_Instruction) {
            return false;
        }

        // 特殊处理Load指令 - 检查是否是从可能变化的内存位置加载
        if (inst instanceof Load_Instruction) {
            if (!isSafeToHoistLoad((Load_Instruction) inst, loop, knownInvariants)) {
                return false;
            }
        }

        //  特殊处理GetElementPtr指令 - 检查索引是否依赖循环变量
        if (inst instanceof Pointer_Instruction) {
            if (!isSafeToHoistGEP((Pointer_Instruction) inst, loop, knownInvariants)) {
                return false;
            }
        }

        // 检查所有操作数
        for (int i = 1; i < inst.get_op().size(); i++) { // 跳过结果操作数
            IR_Value operand = inst.get_op().get(i);

            if (operand instanceof IR_Int_Value || operand instanceof IR_Float_Value) {
                // 常数是不变的
                continue;
            }

            if (operand instanceof IR_Local_Value || operand instanceof IR_VirtualReg_Value) {
                //  检查是否是循环归纳变量（phi指令的结果）
                if (isLoopInductionVariable(operand, loop)) {
                    return false;
                }

                // 检查定义是否在循环外或者是已知的不变式
                Instruction defInst = findDefiningInstruction(operand, loop);
                if (defInst == null) {
                    // 在循环外定义，是不变的
                    continue;
                }

                if (knownInvariants.contains(defInst)) {
                    // 已知的不变式
                    continue;
                }

                // 在循环内定义且不是不变式
                return false;
            }
        }

        // 额外检查指令是否依赖于循环变量
        if (dependsOnLoopVariables(inst, loop)) {
            return false;
        }

        return true;
    }


    /**
     * 检查Load指令是否安全外提
     */
    private boolean isSafeToHoistLoad(Load_Instruction loadInst, Loop loop, Set<Instruction> knownInvariants) {
        IR_Value address = loadInst.get_op().get(1); // Load指令的地址操作数

        // 检查地址计算是否依赖循环变量
        if (address instanceof IR_Local_Value || address instanceof IR_VirtualReg_Value) {
            // 如果地址是由循环内的指令计算得出的，需要检查该指令
            Instruction defInst = findDefiningInstruction(address, loop);
            if (defInst != null) {
                // 地址在循环内计算
                if (defInst instanceof Pointer_Instruction) {
                    // 如果是GEP指令，需要检查其索引是否依赖循环变量
                    if (!isSafeToHoistGEP((Pointer_Instruction) defInst, loop, knownInvariants)) {
                        return false;
                    }
                } else if (!knownInvariants.contains(defInst)) {
                    // 地址计算不是不变式
                    return false;
                }
            }
        }

        //  检查循环中是否有Store指令可能修改此内存位置
        return !mayBeModifiedInLoop(address, loop);
    }

    /**
     *检查GetElementPtr指令是否安全外提
     */
    private boolean isSafeToHoistGEP(Pointer_Instruction gepInst, Loop loop, Set<Instruction> knownInvariants) {
        // 检查所有索引是否为循环不变式
        for (int i = 2; i < gepInst.get_op().size(); i++) { // 跳过结果和基址
            IR_Value index = gepInst.get_op().get(i);

            if (index instanceof IR_Int_Value || index instanceof IR_Float_Value) {
                // 常数索引是安全的
                continue;
            }

            if (index instanceof IR_Local_Value || index instanceof IR_VirtualReg_Value) {
                // 检查是否是循环归纳变量
                if (isLoopInductionVariable(index, loop)) {
                    return false;
                }

                // 检查索引的定义
                Instruction defInst = findDefiningInstruction(index, loop);
                if (defInst != null && !knownInvariants.contains(defInst)) {
                    // 索引在循环内定义且不是不变式
                    return false;
                }
            }
        }

        return true;
    }


    /**
     *检查是否是循环归纳变量
     */
    private boolean isLoopInductionVariable(IR_Value value, Loop loop) {
        // 直接检查是否是phi指令的结果
        for (IR_BaseBlock block : loop.body) {
            for (Instruction inst : block.get_Instruction_List()) {
                if (inst instanceof Phi_Instruction) {
                    if (inst.get_op().size() > 0 &&
                            inst.get_op().get(0) != null &&
                            inst.get_op().get(0).get_ValueName().equals(value.get_ValueName())) {
                        return true;
                    }
                }
            }
        }

        // 间接检查：如果值依赖于phi指令，也认为是归纳变量
        Instruction defInst = findDefiningInstruction(value, loop);
        if (defInst != null) {
            for (int i = 1; i < defInst.get_op().size(); i++) {
                IR_Value operand = defInst.get_op().get(i);
                if (operand instanceof IR_Local_Value || operand instanceof IR_VirtualReg_Value) {
                    if (isPhiResult(operand, loop)) {
                        return true;
                    }
                }
            }
        }
        return false;
    }

    /**
     * 检查值是否是phi指令的结果
     */
    private boolean isPhiResult(IR_Value value, Loop loop) {
        for (IR_BaseBlock block : loop.body) {
            for (Instruction inst : block.get_Instruction_List()) {
                if (inst instanceof Phi_Instruction) {
                    if (inst.get_op().size() > 0 &&
                            inst.get_op().get(0) != null &&
                            inst.get_op().get(0).get_ValueName().equals(value.get_ValueName())) {
                        return true;
                    }
                }
            }
        }
        return false;
    }

    /**
     * 检查指令是否依赖于循环变量
     */
    private boolean dependsOnLoopVariables(Instruction inst, Loop loop) {
        // 获取所有循环归纳变量
        Set<String> inductionVars = getInductionVariables(loop);

        // 递归检查依赖关系
        return hasTransitiveDependency(inst, inductionVars, loop, new HashSet<>());
    }

    /**
     * 获取所有循环归纳变量
     */
    private Set<String> getInductionVariables(Loop loop) {
        Set<String> inductionVars = new HashSet<>();

        for (IR_BaseBlock block : loop.body) {
            for (Instruction inst : block.get_Instruction_List()) {
                if (inst instanceof Phi_Instruction &&
                        inst.get_op().size() > 0 &&
                        inst.get_op().get(0) != null) {
                    inductionVars.add(inst.get_op().get(0).get_ValueName());
                }
            }
        }

        return inductionVars;
    }

    /**
     * 检查传递依赖关系
     */
    private boolean hasTransitiveDependency(Instruction inst, Set<String> inductionVars,
                                            Loop loop, Set<String> visited) {
        for (int i = 1; i < inst.get_op().size(); i++) {
            IR_Value operand = inst.get_op().get(i);

            if (operand instanceof IR_Local_Value) {
                String varName = operand.get_ValueName();

                // 避免无限递归
                if (visited.contains(varName)) {
                    continue;
                }
                visited.add(varName);

                // 直接依赖于归纳变量
                if (inductionVars.contains(varName)) {
                    return true;
                }

                // 检查传递依赖
                Instruction defInst = findDefiningInstruction(operand, loop);
                if (defInst != null && hasTransitiveDependency(defInst, inductionVars, loop, visited)) {
                    return true;
                }

                visited.remove(varName);
            }
        }

        return false;
    }


    /**
     * 找到定义指令
     */
    private Instruction findDefiningInstruction(IR_Value value, Loop loop) {
        for (IR_BaseBlock block : loop.body) {
            for (Instruction inst : block.get_Instruction_List()) {
                if (inst.get_op().size() > 0 &&
                        inst.get_op().get(0) != null &&
                        inst.get_op().get(0).get_ValueName().equals(value.get_ValueName())) {
                    return inst;
                }
            }
        }
        return null; // 在循环外定义
    }

    /**
     * 外提循环不变式
     */
    private void hoistInvariants(Loop loop) {
        // 按照依赖关系排序指令
        List<Instruction> sortedInvariants = topologicalSort(loop.invariants, loop);

        for (Instruction inst : sortedInvariants) {
            // 检查是否安全移动
            if (isSafeToHoist(inst, loop)) {
                // 从原位置移除
                IR_BaseBlock originalBlock = findContainingBlock(inst, loop);
                if (originalBlock != null) {
                    originalBlock.get_Instruction_List().remove(inst);
                }

                // 添加到预头部（保持正确的依赖顺序）
                insertInPreHeader(inst, loop.preHeader);

                // 更新指令所属的基本块
                inst.set_bb(loop.preHeader);
            }
        }
    }

    /**
     *  拓扑排序保证依赖顺序
     */
    private List<Instruction> topologicalSort(Set<Instruction> invariants, Loop loop) {
        List<Instruction> result = new ArrayList<>();
        Set<Instruction> visited = new HashSet<>();
        Set<Instruction> visiting = new HashSet<>();

        for (Instruction inst : invariants) {
            if (!visited.contains(inst)) {
                topologicalSortUtil(inst, invariants, loop, visited, visiting, result);
            }
        }

        return result;
    }

    /**
     * 拓扑排序辅助方法
     */
    private void topologicalSortUtil(Instruction inst, Set<Instruction> invariants, Loop loop,
                                     Set<Instruction> visited, Set<Instruction> visiting,
                                     List<Instruction> result) {
        if (visiting.contains(inst)) {
            // 检测到循环依赖，这种情况下保持原有顺序
            return;
        }

        if (visited.contains(inst)) {
            return;
        }

        visiting.add(inst);

        // 找到所有依赖的指令
        for (int i = 1; i < inst.get_op().size(); i++) {
            IR_Value operand = inst.get_op().get(i);
            if (operand instanceof IR_Local_Value) {
                Instruction defInst = findDefiningInstruction(operand, loop);
                if (defInst != null && invariants.contains(defInst)) {
                    topologicalSortUtil(defInst, invariants, loop, visited, visiting, result);
                }
            }
        }

        visiting.remove(inst);
        visited.add(inst);
        result.add(inst);
    }

    /**
     *  在预头部正确位置插入指令
     */
    private void insertInPreHeader(Instruction inst, IR_BaseBlock preHeader) {
        List<Instruction> preHeaderInsts = preHeader.get_Instruction_List();

        // 找到最后一个分支指令之前的位置
        int insertPos = preHeaderInsts.size();
        for (int i = preHeaderInsts.size() - 1; i >= 0; i--) {
            if (preHeaderInsts.get(i) instanceof Branch_Instruction) {
                insertPos = i;
                break;
            }
        }

        preHeaderInsts.add(insertPos, inst);
    }

    /**
     * 检查是否安全移动
     */
    private boolean isSafeToHoist(Instruction inst, Loop loop) {
        // 检查是否有副作用
        if (inst instanceof Store_Instruction ||
                inst instanceof Call_Instruction) {
            return false;
        }

        // 检查是否支配所有使用点
        IR_Value def = inst.get_op().get(0);
        for (IR_BaseBlock block : loop.body) {
            for (Instruction use : block.get_Instruction_List()) {
                if (usesValue(use, def)) {
                    // 检查定义是否支配使用
                    if (!dominates(findContainingBlock(inst, loop), block)) {
                        return false;
                    }
                }
            }
        }

        return true;
    }

    /**
     * 检查指令是否使用某个值
     */
    private boolean usesValue(Instruction inst, IR_Value value) {
        for (int i = 1; i < inst.get_op().size(); i++) {
            if (inst.get_op().get(i) != null &&
                    inst.get_op().get(i).get_ValueName().equals(value.get_ValueName())) {
                return true;
            }
        }
        return false;
    }

    /**
     * 检查块A是否支配块B
     */
    private boolean dominates(IR_BaseBlock a, IR_BaseBlock b) {
        return dominators.get(b).contains(a);
    }

    /**
     * 找到包含指令的基本块
     */
    private IR_BaseBlock findContainingBlock(Instruction inst, Loop loop) {
        for (IR_BaseBlock block : loop.body) {
            if (block.get_Instruction_List().contains(inst)) {
                return block;
            }
        }
        return null;
    }

    /**
     * 更新分支指令的目标
     */
    private void updateBranchTarget(IR_BaseBlock block, IR_BaseBlock oldTarget, IR_BaseBlock newTarget) {
        List<Instruction> instructions = block.get_Instruction_List();
        if (instructions.isEmpty()) return;

        Instruction lastInst = instructions.get(instructions.size() - 1);
        if (lastInst instanceof Branch_Instruction) {
            Branch_Instruction branch = (Branch_Instruction) lastInst;

            if (branch.br_type()) {
                // 条件分支
                if (branch.get_block1() == oldTarget) {
                    branch.setBlock1(newTarget);
                }
                if (branch.get_block2() == oldTarget) {
                    branch.setBlock2(newTarget);
                }
            } else {
                // 无条件分支
                if (branch.get_block1() == oldTarget) {
                    branch.setBlock1(newTarget);
                }
            }
        }

    }

    private void updatePhiTarget(IR_BaseBlock block, IR_BaseBlock oldTarget, IR_BaseBlock newTarget){
        for(Instruction inst : oldTarget.get_Instruction_List()) {
            if (inst instanceof Phi_Instruction phiInstruction) {
                if(phiInstruction.get_ComingValues().containsKey(block)){
                    IR_Value value = phiInstruction.get_ComingValues().get(block);
                    phiInstruction.remove_ComingValues(block);
                    phiInstruction.set_ComingValue(value,newTarget);
                }
            }
        }
    }
    /**
     * 获取检测到的循环信息
     */
    public List<Loop> getLoops() {
        return loops;
    }

    /**
     * 获取检测到的回边信息
     */
    public List<BackEdge> getBackEdges() {
        return backEdges;
    }

    /**
     * 打印循环信息（用于调试）
     */
    public void printLoopInfo() {
        System.out.println("检测到的回边:");
        for (BackEdge edge : backEdges) {
            System.out.println("  " + edge.from.get_ValueName() + " -> " + edge.to.get_ValueName());
        }

        System.out.println("\n检测到的循环:");
        for (int i = 0; i < loops.size(); i++) {
            Loop loop = loops.get(i);
            System.out.println("  循环 " + i + ":");
            System.out.println("    头部: " + loop.header.get_ValueName());
            System.out.print("    循环体: ");
            for (IR_BaseBlock block : loop.body) {
                System.out.print(block.get_ValueName() + " ");
            }
            System.out.println();
            System.out.println("    不变式数量: " + loop.invariants.size());
        }
    }

    /**
     *检查内存位置是否可能在循环中被修改
     */
    private boolean mayBeModifiedInLoop(IR_Value address, Loop loop) {
        // 简化的别名分析：检查循环中是否有Store指令
        for (IR_BaseBlock block : loop.body) {
            for (Instruction inst : block.get_Instruction_List()) {
                if (inst instanceof Store_Instruction) {
                    Store_Instruction storeInst = (Store_Instruction) inst;
                    IR_Value storeAddress = storeInst.get_op().get(2); // Store指令的地址操作数

                    // 简单的别名检查：如果地址相同或可能重叠，则认为可能被修改
                    if (mayAlias(address, storeAddress)) {
                        return true;
                    }
                }

                // 检查函数调用，函数调用可能修改任何内存
                if (inst instanceof Call_Instruction) {
                    return true; // 保守处理：假设所有函数调用都可能修改内存
                }
            }
        }

        return false;
    }

    /**
     *  简化的别名分析
     */
    private boolean mayAlias(IR_Value addr1, IR_Value addr2) {
        // 非常简化的别名分析
        if (addr1.get_ValueName().equals(addr2.get_ValueName())) {
            return true; // 相同的值一定别名
        }

        // 对于更复杂的情况，这里可以实现更精确的别名分析
        // 现在保守处理：如果不能确定，假设可能别名

        // 检查是否都是基于同一个基址的GEP指令
        Instruction def1 = findDefiningInstructionInLoop(addr1);
        Instruction def2 = findDefiningInstructionInLoop(addr2);

        if (def1 instanceof Pointer_Instruction && def2 instanceof Pointer_Instruction) {
            Pointer_Instruction gep1 = (Pointer_Instruction) def1;
            Pointer_Instruction gep2 = (Pointer_Instruction) def2;

            // 如果基址相同，则可能别名
            if (gep1.get_op().get(1).get_ValueName().equals(gep2.get_op().get(1).get_ValueName())) {
                return true;
            }
        }

        // 保守处理
        return true;
    }

    /**
     * 在整个函数中找到定义指令
     */
    private Instruction findDefiningInstructionInLoop(IR_Value value) {
        // 这里需要在整个函数范围内查找，不仅仅是循环体
        // 实现细节依赖于具体的IR结构
        for(IR_BaseBlock block:curfuncBlock.get_bb_List()){
            for(Instruction inst:block.get_Instruction_List()){
                if(inst instanceof Branch_Instruction || inst instanceof Ret_Instruction || inst instanceof Store_Instruction){
                    continue;
                }
                if(inst.get_op_index(0).equals(value)){
                    return inst;
                }
            }
        }
        return null;
    }

}
